import re
import json
import random

from sonetto.agents.base import BaseAgent, ContextLevel
from sonetto.agents.chat import ChatAgent
from sonetto.agents.tool import ToolAgent
from sonetto.core.message import Message
from sonetto.core.scheduler import Scheduler
from sonetto.core.routing.broadcast import CommonBroadcastConditions

base_url = "https://open.bigmodel.cn/api/paas/v4"
api_key = "ff37fbba1e80d4cb5f38ac819893f4a1.c6T0Zm3Sd9YODRLV"
model = "glm-4-flash"

players = ['玩家1', '玩家2', '玩家3', '玩家4', '玩家5']
roles = ['平民', '平民', '狼人', '预言家', '狼人']

game_rule = '''
你正在参加狼人杀游戏。游戏规则如下：
1. 游戏分为白天和黑夜两个阶段。
2. 白天阶段，所有玩家依次发言，然后投票淘汰一名玩家。
3. 黑夜阶段，狼人可以选择一名玩家投票淘汰，预言家可以查验一名玩家的身份。
4. 游戏中有两大阵营：好人阵营和狼人阵营。
5. 好人阵营包括2名平民和1名预言家，狼人阵营包括2名狼人。但是所有玩家并不知道其他玩家的身份。
5. 对于好人阵营，其目标是找出并投票出狼人，确保好人阵营的胜利。
6. 对于狼人阵营，其目标是淘汰所有好人阵营的玩家，确保狼人阵营的胜利。
7. 游戏中有一个预言家，每晚可以查验一个玩家的身份。
8. 游戏中有一个主持人，负责引导游戏进程。
9. 白天时，你的发言和投票结果会被广播给所有玩家，黑夜时，你的发言只会被主持人看到，所以天黑时不用担心暴露身份。
10. 游戏结束条件：当所有狼人被淘汰或者狼人的人数超过好人的人数时，游戏结束。

你的发言的开头必须以“{}发言：”开头，然后接着你要说的内容。
如 “{}发言：我怀疑玩家2是狼人。”
'''


class PCAgent(BaseAgent):

    def __init__(self, agent_id):
        super().__init__(agent_id)
        self.players = players
        self.roles = roles
        self.player_roles = dict(zip(self.players, self.roles))
        self.alive_players = self.players.copy()

        # 为每一个玩家初始化投票结果
        self.vote_result = {player: 0 for player in self.players}

    def vote(self, target):
        if target in self.alive_players:
            self.vote_result[target] += 1

    def kill_by_vote(self):
        # 按照投票结果淘汰玩家，如果平票则不淘汰
        max_votes = max(self.vote_result.values())
        if list(self.vote_result.values()).count(max_votes) == 1:
            target = max(self.vote_result, key=self.vote_result.get)
            self.alive_players.remove(target)
            self.vote_result = {player: 0 for player in self.alive_players}
            return f'{target}被投票淘汰。'
        else:
            return '平票，无人被淘汰。'

    def kill_by_wolf(self):
        # 选择票数最多的玩家淘汰。如果平票则随机选择一名玩家淘汰
        max_votes = max(self.vote_result.values())
        max_players = [
            player for player in self.vote_result
            if self.vote_result[player] == max_votes
        ]
        if len(max_players) == 1:
            target = max_players[0]
        else:
            target = random.choice(max_players)
        self.alive_players.remove(target)
        self.vote_result = {player: 0 for player in self.alive_players}
        return f'{target}被狼人淘汰。'

    def vote_clear(self):
        self.vote_result = {player: 0 for player in self.alive_players}

    def get_alive_players(self):
        return f'现在场上的存活玩家有：{", ".join(self.alive_players)}。'

    def get_role(self, target):
        return f'{target}的身份是：{self.player_roles[target]}'

    def forward(self, message: Message) -> Message:
        sender = message.sender
        content = message.content
        meta = message.metadata
        receiver = message.receiver

        if sender == '__init__':
            meta['time'] = '白天'
            meta['period'] = '发言'
            meta['route'] = '广播'
            meta['status'] = '未查验'
            return Message(content='''
主持人发言：游戏开始！现在是第一天的白天，大家对其它人的身份都一无所知，请尽情地搅浑场上的局势吧！
现在请玩家1发言。
                           ''',
                           sender=self.agent_id,
                           receiver=self.players[0],
                           metadata=meta)
        elif meta['time'] == '白天' and meta['period'] == '发言':
            meta['route'] = '广播'
            # 白天发言阶段
            # 决定下一个发言对象
            if sender != self.alive_players[-1]:
                next_player = self.alive_players[self.alive_players.index(sender) + 1]
                content = f'请{next_player}发言。'
            else:
                meta['period'] = '投票'
                next_player = self.alive_players[0]
                content = self.get_alive_players(
                ) + f'''现在进入投票阶段。请{next_player}投票选择一名你想要淘汰的玩家，并说明理由。
                投票必须以 json 格式发送，格式如下：
                ```json
                {{"target": "玩家1"}}
                ```
                你也可以选择弃权，同样必须以 json 格式发送，格式如下：
                ```json
                {{"target": "弃权"}}
                ```
                '''
            receiver = next_player
        elif meta['time'] == '白天' and meta['period'] == '投票':
            meta['route'] = '广播'
            # 白天投票阶段
            # 正则表达式提取json对象，并解析投票结果
            vote = re.search(r'\{.*\}', content).group()
            vote = json.loads(vote)
            target = vote['target']
            if target == '弃权':
                content = f'{sender}选择弃权。'
            elif target not in self.alive_players:
                content = f'{sender}选择投票淘汰{target}，但该玩家已被淘汰，故视为弃权。'
            else:
                content = f'{sender}选择投票淘汰{target}。'
                self.vote(target)
            if sender != self.alive_players[-1]:
                next_player = self.alive_players[self.alive_players.index(sender) + 1]
                content = self.get_alive_players(
                ) + f'''请{next_player}投票选择一名你想要淘汰的玩家，并说明理由。
                投票必须以 json 格式发送，格式如下：
                ```json
                {{"target": "玩家1"}}
                ```
                你也可以选择弃权，以 json 格式发送，格式如下：
                ```json
                {{"target": "弃权"}}
                ```
                '''
            else:
                meta['time'] = '黑夜'
                result = self.kill_by_vote()
                if '玩家3' not in self.alive_players and '玩家5' not in self.alive_players:
                    content = result + f'\n游戏结束，好人阵营胜利！'
                elif '玩家1' not in self.alive_players and '玩家2' not in self.alive_players and '玩家4' not in self.alive_players:
                    content = result + f'\n游戏结束，狼人阵营胜利！'
                content = result + f'\n天黑请闭眼。天黑时，狼人和预言家的发言只会被主持人看到，无需担心暴露身份。'
                next_player = self.agent_id
                self.vote_clear()
            receiver = next_player
        elif meta['time'] == '黑夜':
            meta['route'] = '点播'
            # 夜晚阶段
            # 预言家查验身份
            if sender == self.agent_id and '玩家4' in self.alive_players:
                next_player = '玩家4'
                content = self.get_alive_players() + f'''
                请{next_player}选择查验一名玩家的身份，并发送 json 格式的查验选择。
                格式如下：
                ```json
                {{"target": "玩家4"}}
                ```
                '''
            elif sender == '玩家4' and meta['status'] != '已查验':
                meta['status'] = '已查验'
                next_player = '玩家4'
                check = re.search(r'\{.*\}', content).group()
                check = json.loads(check)
                target = check['target']
                content = self.get_role(target) + f'\n请{next_player}铭记并闭眼。'
            elif sender == '玩家4' or sender == self.agent_id:
                next_player = '玩家3' if '玩家3' in self.alive_players else '玩家5'
                content = self.get_alive_players() + f'''
                请狼人：{next_player}选择一名玩家投票淘汰，并发送 json 格式的淘汰选择。
                格式如下：
                ```json
                {{"target": "玩家3"}}
                ```
                你也可以选择弃权，格式如下：
                ```json
                {{"target": "弃权"}}
                ```
你应当选择你认为的好人阵营中的一员进行淘汰，并且尽量不要淘汰到狼人，以确保狼人阵营的胜利。
你的投票结果会和你的队友的投票结果合并。如果你们有不同的投票结果，将会从你们的选择中随机选择一名玩家淘汰。
                '''
            elif sender == '玩家3':
                meta['status'] = '未查验'
                vote = re.search(r'\{.*\}', content).group()
                vote = json.loads(vote)
                target = vote['target']
                if target == '弃权':
                    content = f'{sender}选择弃权。'
                elif target not in self.alive_players:
                    content = f'{sender}选择投票淘汰{target}，但该玩家已被淘汰，故视为弃权。'
                else:
                    self.vote(target)
                    next_player = '玩家5' if '玩家5' in self.alive_players else self.agent_id
                    if next_player == self.agent_id:
                        meta['time'] = '白天'
                        meta['period'] = '发言'
                        meta['route'] = '广播'
                        result = self.kill_by_wolf()
                        next_player = self.alive_players[0]
                        content = result + f'\n天亮请睁眼。'
                        self.vote_clear()
                    else:
                        content = self.get_alive_players() + f'''
                    请狼人：{next_player}选择一名玩家投票淘汰，并发送 json 格式的淘汰选择。
                    格式如下：
                    ```json
                    {{"target": "玩家5"}}
                    ```
                    你也可以选择弃权，格式如下：
                    ```json
                    {{"target": "弃权"}}
                    ```
你应当选择你认为的好人阵营中的一员进行淘汰，并且尽量不要淘汰到狼人，以确保狼人阵营的胜利。
你的投票结果会和你的狼人队友的投票结果合并。如果你们有不同的投票结果，将会从你们的选择中随机选择一名玩家淘汰。
                    '''
            elif sender == '玩家5':
                meta['status'] = '未查验'
                meta['router'] = '广播'
                vote = re.search(r'\{.*\}', content).group()
                vote = json.loads(vote)
                target = vote['target']
                if target != '弃权':
                    self.vote(target)
                meta['time'] = '白天'
                meta['period'] = '发言'
                result = self.kill_by_wolf()
                next_player = self.alive_players[0]
                content = result + f'\n天亮请睁眼。'
                self.vote_clear()
            receiver = next_player
        return Message(content='主持人发言：\n' + content,
                       sender=self.agent_id,
                       receiver=receiver,
                       metadata=meta)


scheduler = Scheduler()

host = PCAgent("主持人")
player1 = ChatAgent("玩家1",
                    api_key=api_key,
                    model=model,
                    system_prompt=game_rule.format("玩家1", '玩家1') + '''
你的角色是：平民。除了发言和投票外，你没有任何特殊能力。
你的最终目的是淘汰所有狼人阵营的玩家，确保好人阵营的胜利。
请小心隐藏自己的身份，以便混淆狼人阵营的判断。
''',
                    base_url=base_url)
player2 = ChatAgent("玩家2",
                    api_key=api_key,
                    model=model,
                    system_prompt=game_rule.format("玩家2", '玩家2') + '''
你的角色是：平民。除了发言和投票外，你没有任何特殊能力。
你的最终目的淘汰所有狼人阵营的玩家，确保好人阵营的胜利。
请小心隐藏自己的身份，以便混淆狼人阵营的判断。
''',
                    base_url=base_url)
player3 = ChatAgent("玩家3",
                    api_key=api_key,
                    model=model,
                    system_prompt=game_rule.format("玩家3", '玩家3') + '''
你的角色是：狼人。你的目标是淘汰所有好人阵营（两个平民，一个预言家）的玩家，确保狼人阵营的胜利。你有一个队友，但你不知道他是谁。
请小心隐藏自己的身份，以便混淆好人阵营的判断。
''',
                    base_url=base_url)
player4 = ChatAgent("玩家4",
                    api_key=api_key,
                    model=model,
                    system_prompt=game_rule.format("玩家4", '玩家4') + '''
你的角色是：预言家。每晚你可以查验一个玩家的身份。
你的目标是淘汰所有狼人阵营的玩家，确保好人阵营的胜利。
请小心隐藏自己的身份，以便混淆狼人阵营的判断。
''',
                    base_url=base_url)
player5 = ChatAgent("玩家5",
                    api_key=api_key,
                    model=model,
                    system_prompt=game_rule.format("玩家5", '玩家5') + '''
你的角色是：狼人。你的目标是淘汰所有好人阵营（两个平民，一个预言家）的玩家，确保狼人阵营的胜利。你有一个队友，但你不知道他是谁。
请小心隐藏自己的身份，以便混淆好人阵营的判断。
''',
                    base_url=base_url)

scheduler.register_agent(host)
scheduler.register_agent(player1)
scheduler.register_agent(player2)
scheduler.register_agent(player3)
scheduler.register_agent(player4)
scheduler.register_agent(player5)

scheduler.set_route(player1.agent_id, host.agent_id)
scheduler.set_route(player2.agent_id, host.agent_id)
scheduler.set_route(player3.agent_id, host.agent_id)
scheduler.set_route(player4.agent_id, host.agent_id)
scheduler.set_route(player5.agent_id, host.agent_id)


def callback(message: Message):
    if message.metadata['time'] == '白天':
        message['route'] = '广播'
    elif message.metadata['time'] == '黑夜':
        message['route'] = '点播'


def broadcast_condition(message: Message):
    return message.metadata['route'] == '广播'


# scheduler.set_callback(callback=callback)
for player in players:
    scheduler.add_broadcast_pattern(
        player, None,
        CommonBroadcastConditions.has_metadata_key_value('time', '白天'),
        '白天时默认广播')

scheduler.add_broadcast_pattern(
    host.agent_id, None,
    CommonBroadcastConditions.has_metadata_key_value('route', '广播'))

scheduler.run_stream('游戏开始！', '游戏结束', [host.agent_id], 40, {})

print([json.loads(ele['content']) for ele in player3.get_context(ContextLevel.SESSION)])
print([json.loads(ele['content']) for ele in player4.get_context(ContextLevel.SESSION)])